import random
from fractions import Fraction
'''
函数说明：每个摇臂提前设置好参数：
m:摇臂的序数
'''
def setting_1(m):
    one = {0 : Fraction(1, 3), 1: Fraction(2, 3), 2:1.0}
    two = {0: 0.5, 2: Fraction(3, 4), 4: 1.0}
    three = {0: 0.5, 5:Fraction(7, 8), 15:1.0}
    four = {0:0.75, 8:Fraction(7, 8), 24:1.0}
    five = {0:Fraction(7, 8), 10:Fraction(31, 32), 100:1.0}
    key = [0, 1, 2, 3, 4]
    value = [one, two, three, four, five]
    all_set = dict(zip(key, value))
    num = random.uniform(0.0, 1.0)
    if num<list(all_set[m].values())[0]:
        feedback = list(all_set[m].keys())[0]
    elif num < list(all_set[m].values())[1]:
        feedback = list(all_set[m].keys())[1]
    elif num < list(all_set[m].values())[2]:
        feedback = list(all_set[m].keys())[2]
    return feedback


'''
函数说明：奖赏函数
m:摇臂的序数
n:该摇臂的摇动次数
'''
def back(m):
    feedback = setting_1(m)
    return feedback

'''函数说明：
基于k摇臂赌博机的贪心算法，假设一共有五个摇臂， 提前设置每个摇臂的参数
k：摇臂数
r：奖赏函数
t: 尝试次数
a：探索概率
'''
def greed(t, a):
    r = 0#累积奖赏
    Q = [0, 0, 0, 0, 0]#初始化每个摇臂的平均奖赏
    count = [0, 0, 0, 0, 0]#初始化每个摇臂的选中次数
    for num in range(t):
        if random.uniform(0.0, 1.0) < a or Q == [0, 0, 0, 0, 0]:
            k1 = random.randint(0, 4)
        else:
            k1 = Q.index(max(Q))
        new_v = back(k1)#最新一次摇臂后的奖赏
        r = r + new_v
        Q[k1] = (Q[k1]*count[k1] + new_v)/int(count[k1] + 1)
        count[k1] = count[k1] + 1
    return r, Q, count

if __name__ == '__main__':
    r, q, count = greed(1000, 0.5)
    for i in range(0, 5):
        print('摇臂{}的平均奖赏为{},摇的次数为{}'.format(i+1, q[i], count[i]))
    print('累计奖赏为:' + str(r))


